package com.zxjbyte.yiyi.framework.common.core.okhttp;

import com.zxjbyte.yiyi.framework.common.core.okhttp.callback.IProgressCallback;
import okhttp3.MediaType;
import okhttp3.ResponseBody;
import okio.BufferedSource;
import okio.ForwardingSource;
import okio.Okio;
import okio.Source;
import org.jetbrains.annotations.NotNull;
import org.springframework.lang.Nullable;

import java.io.IOException;

/**
 * 对 ResponseBody 类进行封装，用于在读取服务端返回的数据流的同时，还会记录已读取数据的大小，这样就能够实现下载进度的统计
 * @Author zhangxingjia
 * @Date 2023/5/6 16:08
 * @Version: 1.0
 */
public class CountingResponseBody extends ResponseBody {

    private final ResponseBody responseBody;
    private final IProgressCallback progressCallback;

    private BufferedSource bufferedSource;

    public CountingResponseBody(ResponseBody responseBody, IProgressCallback progressCallback) {
        this.responseBody = responseBody;
        this.progressCallback = progressCallback;
    }

    @Override
    public long contentLength() {
        return responseBody.contentLength();
    }

    @Nullable
    @Override
    public MediaType contentType() {
        return responseBody.contentType();
    }

    @NotNull
    @Override
    public BufferedSource source() {
        if (bufferedSource == null) {
            bufferedSource = Okio.buffer(new CountingSource(responseBody.source()));
        }
        return bufferedSource;
    }

    private class CountingSource extends ForwardingSource {
        private long totalBytesRead = 0L;

        public CountingSource(Source delegate) {
            super(delegate);
        }

        @Override
        public long read(@NotNull okio.Buffer sink, long byteCount) throws IOException {
            long bytesRead = super.read(sink, byteCount);
            if (bytesRead != -1) {
                totalBytesRead += bytesRead;
            }
            // 如果responseBody.contentLength() 的返回值为 -1 表示无法获取响应体的长度，这通常是由于服务器未正确设置 Content-Length 响应头造成的
            progressCallback.onDownloadProgress(totalBytesRead, responseBody.contentLength(), bytesRead == -1);
            return bytesRead;
        }
    }
}
